import random
import numpy as np
import matplotlib.pyplot as plt


def dampnm(fun, gfun, hess, x0):
    # 用牛顿法求解无约束问题
    # x0是初始点，fun，gfun和hess分别是目标函数值，梯度，海森矩阵的函数
    maxk = 500
    rho = 0.55
    sigma = 0.4
    k = 0
    epsilon = 1e-5

    f = open("牛顿.txt", "w")
    W = np.zeros((2, 20000))

    while k < maxk:
        W[:, k] = x0
        gk = gfun(x0)
        Gk = hess(x0)
        dk = -1.0 * np.linalg.solve(Gk, gk)
        print(k, np.linalg.norm(dk))
        f.write(str(k) + '   ' + str(np.linalg.norm(gk)) + "\n")
        if np.linalg.norm(dk) < epsilon:
            break

        x0 += dk
        k += 1
    W = W[:, 0:k + 1]  # 记录迭代点
    return x0, fun(x0), k, W


# 函数表达式fun
fun = lambda x: 100 * (x[0] ** 2 - x[1]) ** 2 + (x[0] - 1) ** 2

# 梯度向量 gfun
gfun = lambda x: np.array([400 * x[0] * (x[0] ** 2 - x[1]) + 2 * (x[0] - 1), -200 * (x[0] ** 2 - x[1])])

# 海森矩阵 hess
hess = lambda x: np.array([[1200 * x[0] ** 2 - 400 * x[1] + 2, -400 * x[0]], [-400 * x[0], 200]])

if __name__ == "__main__":
    X1 = np.arange(-1.5, 1.5 + 0.05, 0.05)
    X2 = np.arange(-3.5, 2 + 0.05, 0.05)
    [x1, x2] = np.meshgrid(X1, X2)
    f = 100 * (x2 - x1 ** 2) ** 2 + (1 - x1) ** 2  # 给定的函数
    plt.contour(x1, x2, f, 40)  # 画出函数的20条轮廓线

    x0 = np.array([-1.2, 1])
    out = dampnm(fun, gfun, hess, x0)
    print(out[2], out[0])

    W = out[3]
    print(W[:, :])

    plt.plot(W[0, :], W[1, :], 'g*-')
    plt.show()
